import torch.nn as nn


"""
AlexNet网络结构
"""
class AlexNet(nn.Module):

    def __init__(self, classes):
        super().__init__()
        self.__create_nn(classes)

    def __create_nn(self, classes):
        self.features = [
            nn.Conv2d(3, 48, (11, 11), (4, 4), 2),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(48, 128, (5, 5), (4, 4), 0),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2),
            nn.Conv2d(128, 192, (3, 3), (1, 1), 1),
            nn.ReLU(True),
            nn.Conv2d(192, 192, (3, 3), (1, 1), 1),
            nn.ReLU(True),
            nn.Conv2d(192, 128, (3, 3), (1, 1), 1),
            nn.ReLU(True),
            nn.MaxPool2d(3, 2)
        ]

        self.full_connect = [
            nn.Flatten(1),
            nn.Dropout(0.5),
            nn.Linear(128 * 6 * 6, 2048),
            nn.ReLU(True),
            nn.Dropout(0.5),
            nn.Linear(2048, 2048),
            nn.ReLU(True),
            nn.Linear(2048, classes)
        ]

        sequential = self.features + self.full_connect
        self.sequential = nn.Sequential(*sequential)
        self.__init_weight()

    def __init_weight(self):
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(module.weight)
                if module.bias is not None:
                    nn.init.constant_(module.bias, 0)
            else:
                nn.init.normal_(module.weight, 0, 0.01)
                nn.init.constant_(module.bias, 0)

    def forward(self, x):
        x = self.sequential(x)
        return x

if __name__ == "__main__":
    net = AlexNet(5)
